import sys

file1=open(sys.argv[1],'r')# gene_gff
file2=open(sys.argv[2],'r')# TE_gff
chr_TE_dict={}
for line in file2:
    if line.startswith('#'):
        continue
    line_list=line.strip().split()
    if '(' in line_list[-1]:
        continue
    if line_list[0] not in chr_TE_dict:
        chr_TE_dict[line_list[0]]={}
    if line_list[-1] not in chr_TE_dict[line_list[0]]:
        chr_TE_dict[line_list[0]][line_list[-1]]=[]
    if int(line_list[1])<int(line_list[2]):
        chr_TE_dict[line_list[0]][line_list[-1]].append([int(line_list[1]),int(line_list[2])])
    else:
        chr_TE_dict[line_list[0]][line_list[-1]].append([int(line_list[2]),int(line_list[1])])
    # print(chr_TE_dict)
    # sys.exit()
    
chr_intege_dict={}
total_length=0
for line in file1:
    line_list=line.strip().split('\t')
    if line_list[0] not in chr_intege_dict.keys():
        chr_intege_dict[line_list[0]]=[]
    if int(line_list[1])<int(line_list[2]):
        chr_intege_dict[line_list[0]].append([int(line_list[1]),int(line_list[2])])
        total_length+=int(line_list[2])-int(line_list[1])
    else:
        chr_intege_dict[line_list[0]].append([int(line_list[2]),int(line_list[1])])
        total_length+=int(line_list[1])-int(line_list[2])
    
file1.close()
file1=open(sys.argv[1],'r')
# print('基因区长度\t'+total_length)      
# length=0
out_dict={'CACTA_TIR_transposon':0,'Copia_LTR_retrotransposon':0,'Gypsy_LTR_retrotransposon':0,'hAT_TIR_transposon':0,'helitron':0,'long_terminal_repeat':0,\
'LTR_retrotransposon':0,'Mutator_TIR_transposon':0,'PIF_Harbinger_TIR_transposon':0,'repeat_region':0,'target_site_duplication':0,'Tc1_Mariner_TIR_transposon':0,'TE':0}

def get_inter_dict(chr_TE_dict,chr_intege_dict,out_dict):
    # out_dict={}
    for key0,value0 in chr_TE_dict.items():
        # print(key)
        try:
            gene_list=sorted(chr_intege_dict[key0],key=lambda x:x[0])
        except KeyError:
            continue
        for key1,value1 in value0.items():
            if key1 not in out_dict:
                continue
                # out_dict[key1]=0
            TE_list=sorted(value1,key=lambda x:x[0])
            index=0
            length=0
            add_length=0
            # while index<=len(TE_list):
            for intro_pos in gene_list:
                for te_pos in TE_list[index:]:
                    # print(intro_pos,te_pos)
                    if te_pos[0]>=intro_pos[0]:
                        if te_pos[0]>intro_pos[1]:
                            break
                        elif te_pos[1]<=intro_pos[1]:
                            length+=te_pos[1]-te_pos[0]+1
                        elif te_pos[1]>=intro_pos[1]:
                            length+=intro_pos[1]-te_pos[0]+1
                    elif te_pos[0]<=intro_pos[0]:
                        if te_pos[1]<=intro_pos[0]:
                            index+=1
                            continue
                        
                        elif te_pos[1]<=intro_pos[1]:
                            length+=te_pos[1]-intro_pos[0]+1
                        elif te_pos[1]>=intro_pos[1]:
                            length+=intro_pos[1]-intro_pos[0]+1
                    # print(length,te_pos,intro_pos)
                    # sys.exit()  
                    
                   
                    
            out_dict[key1]+=length
    return out_dict
out_dict1=get_inter_dict(chr_TE_dict,chr_intege_dict,out_dict)
# print('TE 长度\t'+length)
# print(length)
# print(out_dict1)
file3=open(sys.argv[3]+'.inter.count','w')
file3.write('基因区长度\t'+str(total_length)+'\n')
for key,value in out_dict1.items():
    file3.write(key+'\t'+str(value/total_length)+'\n')
up_chr_intege_dict={}
total_length=0
for line in file1:
    line_list=line.strip().split('\t')
    if line_list[0] not in up_chr_intege_dict.keys():
        up_chr_intege_dict[line_list[0]]=[]
    up_chr_intege_dict[line_list[0]].append([int(line_list[1])-2000,int(line_list[1])])
    up_chr_intege_dict[line_list[0]].append([int(line_list[2]),int(line_list[2])+2000])
    total_length+=4000
out_dict={'CACTA_TIR_transposon':0,'Copia_LTR_retrotransposon':0,'Gypsy_LTR_retrotransposon':0,'hAT_TIR_transposon':0,'helitron':0,'long_terminal_repeat':0,\
'LTR_retrotransposon':0,'Mutator_TIR_transposon':0,'PIF_Harbinger_TIR_transposon':0,'repeat_region':0,'target_site_duplication':0,'Tc1_Mariner_TIR_transposon':0,'TE':0}


out_dict2=get_inter_dict(chr_TE_dict,up_chr_intege_dict,out_dict)
file4=open(sys.argv[3]+'.up.count','w')
file4.write('基因上下游长度\t'+str(total_length)+'\n')
for key,value in out_dict2.items():
    file4.write(key+'\t'+str(value/total_length)+'\n')